import copy
import glob
import os
from os.path import join

import numpy as np
from PIL import Image
from tqdm import tqdm

from models.unet.unet import Unet
from utils.train.imageUtils import letterbox_image
from utils.dataset.patient import PatientData
from utils.evaluation.miou import miou_Unet, compute_mIoU


def compute(net_name, masks, num_classes):
    masksInd = np.arange(len(masks))
    for ind in masksInd:
        masks[ind] = Image.fromarray(masks[0]).convert('L')

    pred_dir = "miou_pr_dir_" + net_name
    pred_imgs = [join(pred_dir, str(x) + ".png") for x in masksInd]

    name_classes = ["background", "target"]
    compute_mIoU(masks, pred_imgs, num_classes, name_classes)


def init(net_name, images, num_classes, model_path, model_image_size):
    unet = miou_Unet(
        net=Unet(num_classes=num_classes, in_channels=model_image_size[-1]).eval(),
        num_classes=num_classes,
        model_path=model_path,
        model_image_size=model_image_size)

    if not os.path.exists("miou_pr_dir_" + net_name):
        os.makedirs("miou_pr_dir_" + net_name)

    for ind in tqdm(np.arange(len(images))):
        image = images[ind]
        old_img = copy.deepcopy(image)
        image, nw, nh = letterbox_image(image=image, size=(model_image_size[1], model_image_size[0]))
        image = [np.array(image) / 255]
        image = np.transpose(image, (0, 3, 1, 2))
        image = unet.detect_image(old_img=old_img, images=image, nw=nw, nh=nh)
        image.save("miou_pr_dir_" + net_name + "/" + str(ind) + ".png")


def main():
    netName = "Unet"
    model_path = "checkpoint_unet_1/Epoch95-Total_Loss0.4983-Val_Loss0.4981.pth"
    data_dir = "../../../data/RVSC2012/Test1Set"
    model_image_size = (256, 256, 3)
    num_classes = 3
    maskType = 'both'
    assert num_classes in ['inner', 'outer', 'both']

    glob_search = os.path.join(data_dir, "patient*")
    patient_dirs = sorted(glob.glob(glob_search))
    images = []
    inner_masks = []
    outer_masks = []
    for patient_dir in patient_dirs:
        p = PatientData(patient_dir)
        for image in p.images:
            images.append(Image.fromarray(image.astype(np.uint8)))
        inner_masks += p.endocardium_masks
        outer_masks += p.epicardium_masks

    if maskType == 'both':
        masks = np.asarray(inner_masks) + np.asarray(outer_masks)
    elif maskType == 'inner':
        masks = np.asarray(inner_masks)
    else:
        masks = np.asarray(outer_masks)

    init(net_name=netName,
         images=images,
         num_classes=num_classes,
         model_path=model_path,
         model_image_size=model_image_size)
    compute(net_name=netName, masks=masks, num_classes=num_classes)


if __name__ == "__main__":
    main()
